#include <string>
#include <algorithm>
#include <iostream>
#include <stdexcept>
#include <vector>
#include <deque>
#include <sstream>
#include <fstream>
#include <boost/array.hpp>
#include <boost/filesystem/path.hpp>
#include <boost/filesystem/operations.hpp>
#define DEBUG_LEVEL 6
#include "components/auto_restore_dir.h"
#include "components/portable_timer.h"
#include "components/debug_func.h"
#include "components/codeml.h"
#include "components/perlutil.h"
#include "components/progress_indicator.h"
#include "components/ultoa.h"
#define _MY_STL_MACROS
#include "components/stlhelper.h"
#include "get_arg.h"
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <ext/hash_map>
using std::cerr;
using boost::array;
using boost::filesystem::path;
using std::fill;
using std::cin;
using std::cout;
using std::deque;
using std::string;
using std::vector;
using std::sort;
using std::replace;
using std::ofstream;
using std::ifstream;



//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//		Main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888










//________________________________________________________________________________________

//	remove_stop_codons

//________________________________________________________________________________________
unsigned char ACTG [256];
unsigned char test_stops [512];
inline unsigned codon_to_512(char a, char b, char c)
{
	return (ACTG[a] << 6) | (ACTG[b] << 3) | (ACTG[c]);
}

void remove_stop_codons(deque<string>& sequences, deque<string>& names, ostream& err_strm)
{
	fill (ACTG, ACTG + 256, 4);
	ACTG['A'] = ACTG['a'] = 0;
	ACTG['C'] = ACTG['c'] = 1;
	ACTG['T'] = ACTG['t'] = 2;
	ACTG['G'] = ACTG['g'] = 3;
	fill(test_stops,  test_stops + 512,  0);
	test_stops[codon_to_512('T', 'A', 'A')] = 1;
	test_stops[codon_to_512('T', 'G', 'A')] = 1;
	test_stops[codon_to_512('T', 'A', 'G')] = 1;

	for (unsigned i = 0; i < sequences.size(); ++i)
	{
		bool replaced_stop = false;
		unsigned stop_pos = 0;
		for (unsigned j = 0; j < sequences[i].length(); j+=3)
		{
			if (test_stops[codon_to_512(sequences[i][j], sequences[i][j + 1], sequences[i][j + 2])])
			{
				sequences[i][j]		=
				sequences[i][j + 1]  =
				sequences[i][j + 2]  =	'-';
				stop_pos = j / 3 + 1;
				replaced_stop = true;
			}
		}
		if (replaced_stop)
		{
			err_strm << "Replaced stop codon for " << names[i] << " at position " <<
						stop_pos << " (counting from 1) \n";
		}
	}

}




//________________________________________________________________________________________

//	analyse_all_orthologs

//		die if return false

//________________________________________________________________________________________


bool analyse_all_orthologs(	deque<string>& names,
							deque<string>& sequences,
							ostream& err_strm,
							const string& kaks_exe_path,
							vector<t_codeml_results>& results,
							bool codeml_not_yn00)

{
	
	unsigned cnt_sequences = sequences.size();
	cerr << "\t" << cnt_sequences << "\tsequences altogether:\n";
	unsigned cnt_pairs = cnt_sequences * (cnt_sequences-1) / 2;
	cerr << "\t" << cnt_pairs << "\talignments to be processed...\n";

	
	// time whole process
	t_portable_timer timer;
	t_portable_timer timer_initial;
	{
		// get data in batches
		unsigned batchsz = std::min(cnt_pairs / 50, 400U)+1;
		t_progress_indicator dots(cerr, batchsz);
		for (unsigned i = 0; i < cnt_sequences; ++i)
		{
			for (unsigned ii = 0; ii < cnt_sequences; ++ii)
			{
				// skip if below diagonal
				if (ii <= i)
					continue;
				vector<string> seq_pair;
				seq_pair.push_back(sequences[i]);
				seq_pair.push_back(sequences[ii]);
				vector<string> name_pair;
				name_pair.push_back(names[i]);
				name_pair.push_back(names[ii]);
				kaks_analyse (	seq_pair, name_pair, results, kaks_exe_path, codeml_not_yn00);
				if (dots.value() == batchsz-1)
				{
					cerr << "\tThe entire job will take around " <<
							t_portable_timer::pretty_time(
											timer_initial.reset() * cnt_pairs / batchsz);
					unsigned cnt_dots = cnt_pairs / batchsz;
					if (cnt_dots > 150)
						cerr << "\t(i.e. " << (cnt_dots + 49) / 50.0 << " lines of dots!)\n";
					else
                        cerr << "\t(i.e. " << cnt_dots << "  dots.)\n";
				}
				dots++;
			}
		}
	}
	cerr << "\tAnalysing orthologs has taken\t" << timer.pretty_reset() << "\n";
	return true;
}











//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	main


//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
bool read_sequences(std::deque<string>& names, std::deque<string>& sequences)
{
	unsigned cnt_line = 0;
	unsigned seq_len = 0;
	while (cin.good())
	{
		++cnt_line;
		string name;
		getline(cin, name,  '\t');
		string seq;
		getline(cin, seq,  '\n');
		if (cin)
		{
			if (!seq_len)
				seq_len = seq.length();
			if (seq_len != seq.length())
			{
				cerr << "Error!!\n\tSequence length changed on line " << cnt_line <<
						" from " << seq_len << " to " << seq.length() << std::endl;
				return false;
			}
			sequences.push_back(seq);
			names.push_back(name);
		}
	}
	return true;

}

bool output_results(const string& path_prefix, vector<t_codeml_results>& results)
{
	array<string, 9> categories = {{"ka", "ks", "kaks", "kappa", "dT", "count_a", "count_s", "lnL"}};
	for (unsigned i = 0; i < categories.size(); ++i)
	{
		ofstream res_strm ((path_prefix + categories[i] + ".data").c_str());
		if (!res_strm)
		{
			cerr << "\n\nError\n\tUnable to open " << path_prefix + categories[i] + ".data\n";
			return false;
		}
		for (unsigned j = 0; j < results.size(); ++j)
		{
			if (results[j].error_str.length())
				continue;
			res_strm << results[j].name1 << "\t" << results[j].name2 << "\t";
			switch (i)
			{
				case 0: res_strm << results[j].dN		<< "\n"; break;
				case 1: res_strm << results[j].dS		<< "\n"; break;
				case 2: res_strm << results[j].dN_dS	<< "\n"; break;
				case 3: res_strm << results[j].kappa	<< "\n"; break;
				case 4: res_strm << results[j].dT 		<< "\n"; break;
				case 5: res_strm << results[j].N 		<< "\n"; break;
				case 6: res_strm << results[j].S 		<< "\n"; break;
				case 7: res_strm << results[j].lnL 		<< "\n"; break;
			}
		}
	}

	ofstream res_strm ((path_prefix + ".errors").c_str());
	for (unsigned j = 0; j < results.size(); ++j)
		if (results[j].error_str.length())
		{
				res_strm	<< results[j].name1 << "\t" << results[j].name2 << "\t"
							<< results[j].error_str << "\n";
			continue;
		}

	return true;
}

int main(int argc, char* argv[])
{
	// don't use C streams
	std::ios::sync_with_stdio(false);
	
	
	// get command line arguments

	// if no arguments print usage and exit
	if (argc == 1)
	{
		char * my_argv[2] ={ argv[0], "-h"};
		int argc = 2;
		t_program_args args;
		if (!process_command_line_options(argc, my_argv, args))
			return 1;
	}
	t_program_args args;
	if (!process_command_line_options(argc, argv, args))
		return 1;
	

	cerr << "\tPAML control file from " << args.paml << " kaks.ctl\n";
	
	char curr_path_str [1024];
	string tmp_file_path = args.tempdir.length() ?
									args.tempdir :
									path(getcwd(curr_path_str,
												sizeof(curr_path_str) /
												sizeof(char)))
												.branch_path().string() + "/kaks_tmp";
	
	
	
	//
	// read in sequences
	//
	std::deque<string> sequences;
	std::deque<string> names;
	if (!read_sequences(names, sequences))
		return 1;
	
	//
	// codeml cannot handle stops
	//
	ofstream err_strm (args.error_file.c_str());
	if (!err_strm)
	{
		cerr << "\n\nError\n\tUnable to open " << args.error_file << "\n";
		return 1;
	}
	remove_stop_codons(sequences, names, err_strm);

	
	vector<t_codeml_results> results;
	
	// use kaks_tmp as temp directory if not specified
	{								
		// save original directory for restore later
		t_auto_restore_dir save_directory;
		
		
		
		bool remove_tmp_file = false;
		if(!prepare_temp_directory(tmp_file_path,
									args.fixed_kappa, args.kappa,
									args.codeml, remove_tmp_file))
			return false;
		
		if (!analyse_all_orthologs(names, sequences, err_strm,
							args.paml + " kaks.ctl", results, args.codeml))
			return 1;
		if (args.fixed_kappa)
			for (unsigned i = 0; i < results.size(); ++i)
				results[i].kappa = args.kappa;
	}	
		

	if (!output_results(args.output_prefix, results))
		return 1;


	return 0;
}


